# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: jaewon

import datetime
import logging
import math
import pprint
import sys

import tabulate

from coin.base.datetime_util import to_datetime
from coin.base.param_util import to_list
from coin.base.timestamp import get_timestamp
from coin.exchange.base.order_gateway import OrderSide, OrderType
from coin.exchange.okex_futures.types import OkexFuturesOrderSide
from coin.strategy.mm.calculator import MovingSum
from coin.strategy.mm.bitmex.lazy_pull_mm_executor import BitmexLazyPullMMExecutor
from coin.strategy.mm.okex_futures.simple_mm_executor_2 import OkexFuturesSimpleMMExecutor_v2
from coin.strategy.mm.base.okex_futures_bitmex_strategy_base import (BitmexFuturesProduct,
                                                                     OkexCurrency,
                                                                     OkexFuturesProduct,
                                                                     StrategyBase)


def ceil_price(price, tick_size):
  return math.ceil(price / tick_size) * tick_size


def floor_price(price, tick_size):
  return math.floor(price / tick_size) * tick_size


def norm_float(value, precision):
  return float(('%%.%df' % precision) % value)


def calculate_unrealized_pnl(pos, avg_price, ask0_price, bid0_price):
  if pos == 0:
    return 0
  mid_price = (ask0_price + bid0_price) / 2.
  return pos * (1. / avg_price - 1. / mid_price)


def get_side_sign(order_side, okex_futures=False):
  if okex_futures:
    if (order_side == OkexFuturesOrderSide.BUY_OPEN
        or order_side == OkexFuturesOrderSide.BUY_CLOSE):
      return 1
    elif (order_side == OkexFuturesOrderSide.SELL_OPEN
          or order_side == OkexFuturesOrderSide.SELL_CLOSE):
      return -1
  else:
    if order_side == OrderSide.BUY:
      return 1
    elif order_side == OrderSide.SELL:
      return -1
  raise ValueError('Invalid order side: %s' % order_side)


class PassivePassiveStrategy(StrategyBase):
  def __init__(self, strategy_params, dry_run=True):
    super().__init__(strategy_params['product_1'],
                     strategy_params['product_2'], ['BTC'],
                     logger=logging.getLogger('PassivePassiveStrategy'))
    self._strategy_params = strategy_params
    self._dry_run = dry_run

    self._product_1 = strategy_params['product_1']
    self._product_2 = strategy_params['product_2']

    # Pricing parameters
    self._basis_ma_window = strategy_params['basis_ma_window_min'] * 60 * (10**9)
    self._true_book_1 = strategy_params['true_book_1']  # OkexFutures
    self._true_book_2 = strategy_params['true_book_2']  # Bitmex
    self._edge_1 = strategy_params['edge_bp_1'] / 10000.
    self._edge_2 = strategy_params['edge_bp_2'] / 10000.

    # Execution parameters
    self._lot_size_1 = strategy_params['lot_size_1']
    self._lot_size_2 = strategy_params['lot_size_2']
    self._reserve = strategy_params['reserve']
    self._imbalance_allowance = strategy_params['imbalance_allowance']

    self._bs_pos = 0.
    self._sb_pos = 0.
    self._mid_basis_ma = MovingSum(self._basis_ma_window)

    self._executor_1 = None
    self._executor_2 = None

    self._last_feature_dump_ts = 0

    self._logger.info('Strategy params:\n%s', pprint.pformat(strategy_params))

    # Print header
    self._logger.info('\nheader_fill_dump,time,time_hr,exchange,side,price,qty,'
                      'pos_1,pos_2,bs_pos,sb_pos,balance_1,balance_2')
    self._logger.info('\nheader_feature_dump,time,time_hr,'
                      'ask0_1,bid0_1,askt_1,bidt_1,midt_1,'
                      'ask0_2,bid0_2,askt_2,bidt_2,midt_2,'
                      'mid_basis_ma,askp_1,bidp_1,askp_2,bidp_2'
                      'pos_1,pos_2,bs_pos,sb_pos,balance_1,balance_2')

  def on_base_init(self, ioloop, queue, flow_sub):
    super().on_base_init(ioloop, queue, flow_sub)

  def _get_pos_balance_1(self):
    pos_from_og_1 = self._og_1.get_okex_futures_position(self._product_1)
    long_pos_1 = pos_from_og_1.long * 100.
    short_pos_1 = -pos_from_og_1.short * 100.
    long_unrealized_pnl_1 = calculate_unrealized_pnl(long_pos_1,
                                                     pos_from_og_1.long_avg_price,
                                                     self._book_1.ask0().price,
                                                     self._book_1.bid0().price)
    short_unrealized_pnl_1 = calculate_unrealized_pnl(short_pos_1,
                                                      pos_from_og_1.short_avg_price,
                                                      self._book_1.ask0().price,
                                                      self._book_1.bid0().price)
    unrealized_pnl_1 = long_unrealized_pnl_1 + short_unrealized_pnl_1

    pos_1 = long_pos_1 + short_pos_1
    balance_1 = self._og_1.get_balance('BTC') + unrealized_pnl_1
    return pos_1, balance_1

  def _get_pos_balance_2(self):
    pos_2 = self._og_2.get_position(self._product_2)
    avg_entry_price_2 = self._og_2.get_average_entry_price(self._product_2)
    unrealized_pnl_2 = calculate_unrealized_pnl(pos_2,
                                                avg_entry_price_2,
                                                self._book_2.ask0().price,
                                                self._book_2.bid0().price)
    balance_2 = self._og_2.get_balance('BTC') + unrealized_pnl_2
    return pos_2, balance_2

  def _print_balance(self):  # Override the original version.
    if self._book_1 is None or self._book_2 is None:
      return

    pos_1, balance_1 = self._get_pos_balance_1()
    pos_2, balance_2 = self._get_pos_balance_2()

    table = [['*', 'BTC', balance_1 + balance_2, balance_1, balance_2],
             ['*', 'USD', pos_1 + pos_2, pos_1, pos_2], ['*', 'USD_BS', self._bs_pos, 0,
                                                         0], ['*', 'USD_SB', self._sb_pos, 0, 0]]

    self._logger.info(
        '----- Balance:\n* %s\n%s' %
        (str(datetime.datetime.now()),
         tabulate.tabulate(
             table, headers=['*', 'Currency', 'Total', 'OKEX Futures', 'Bitmex'], floatfmt=".6f")))

  def init_executor(self):
    self._logger.info('Intializing passive executor...')

    executor_params_1 = {
        'product': self._product_1,
        'lot_size': self._lot_size_1 / 100.,  # Per contract
        'stickiness': 0.,
        'min_pos': -self._reserve / 100.,
        'max_pos': self._reserve / 100.,
        'order_update_period': 2.0,  # Second
    }
    self._executor_1 = OkexFuturesSimpleMMExecutor_v2(executor_params_1, self._og_1)

    executor_params_2 = {
        'product': self._product_2,
        'lot_size': self._lot_size_2,
        'min_pos': -self._reserve,
        'max_pos': self._reserve,
        'order_update_period': 2.0,  # Second
        'post_only': True,
        'pull_threshold_bp': 2.5
    }
    self._executor_2 = BitmexLazyPullMMExecutor(executor_params_2, self._og_2)

    # Register fill callback
    self._og_1.register_fill_callback(self.on_fill_1)
    self._og_2.register_fill_callback(self.on_fill_2)

    # Reset bs_pos, sb_pos
    net_pos = (self._og_1.get_position(self._product_1) * 100.
               + self._og_2.get_position(self._product_2))
    self._bs_pos = net_pos / 2.
    self._sb_pos = net_pos / 2.

    self._logger.info('Initial BS=%.2f, SB=%.2f', self._bs_pos, self._sb_pos)

  def on_fill_1(self, fill_average_price, fill_qty, order):
    if (order.order_side == OkexFuturesOrderSide.BUY_OPEN
        or order.order_side == OkexFuturesOrderSide.BUY_CLOSE):
      self._bs_pos += fill_qty * 100.
    elif (order.order_side == OkexFuturesOrderSide.SELL_OPEN
          or order.order_side == OkexFuturesOrderSide.SELL_CLOSE):
      self._sb_pos -= fill_qty * 100.
    else:
      raise ValueError(order)
    self.dump_fill(1, fill_average_price, fill_qty * 100., order.order_side)

  def on_fill_2(self, fill_price, fill_qty, order):
    if order.order_side == OrderSide.BUY:
      self._sb_pos += fill_qty
    elif order.order_side == OrderSide.SELL:
      self._bs_pos -= fill_qty
    else:
      raise ValueError(order)
    self.dump_fill(2, fill_price, fill_qty, order.order_side)

  def dump_fill(self, exchange_id, fill_price, fill_qty, order_side, timestamp=None):
    # Fill header:
    # "FILL_DUMP",time,time_hr,exchange,side,price,qty,
    # pos_1,pos_2,bs_pos,sb_pos,balance_1,balance_2
    timestamp = timestamp or get_timestamp()
    pos_1, balance_1 = self._get_pos_balance_1()
    pos_2, balance_2 = self._get_pos_balance_2()
    side_sign = get_side_sign(order_side, exchange_id == 1)
    self._logger.info('\n"FILL_DUMP",%d,"%s",%d,%d,%f,%f,'
                      '%f,%f,%f,%f,%f,%f',
                      timestamp,
                      to_datetime(timestamp),
                      exchange_id,
                      side_sign,
                      fill_price,
                      fill_qty,
                      pos_1,
                      pos_2,
                      self._bs_pos,
                      self._sb_pos,
                      balance_1,
                      balance_2)

  def on_book(self):
    if self._book_1 is None or self._book_2 is None:
      return

    # Check book freshness
    book_ts = max(self._book_1.timestamp, self._book_2.timestamp)
    min_ts = min(self._book_1.timestamp, self._book_2.timestamp)
    if book_ts - min_ts > (2 * 10**9):
      return

    # Check book sanity
    if (self._book_1.ask0().price <= self._book_1.bid0().price
        or self._book_2.ask0().price <= self._book_2.bid0().price):
      return

    ask0_1 = self._book_1.ask0().price
    bid0_1 = self._book_1.bid0().price
    askt_1 = self._book_1.get_notional_asks_by_qty(self._true_book_1 / 100.)[1][-1][0]
    bidt_1 = self._book_1.get_notional_bids_by_qty(self._true_book_1 / 100.)[1][-1][0]
    midt_1 = (askt_1 + bidt_1) / 2.

    ask0_2 = self._book_2.ask0().price
    bid0_2 = self._book_2.bid0().price
    askt_2 = self._book_2.get_notional_asks_by_qty(self._true_book_2)[1][-1][0]
    bidt_2 = self._book_2.get_notional_bids_by_qty(self._true_book_2)[1][-1][0]
    midt_2 = (askt_2 + bidt_2) / 2.

    mid_diff = midt_1 - midt_2
    self._mid_basis_ma.update(book_ts, mid_diff)

    if len(self._mid_basis_ma) < 2000:
      return

    # TODO(jaewon): Parameterize
    t1 = 0.01
    t2 = 0.5

    midt_2_proj_1 = midt_2 + float(self._mid_basis_ma.average)
    askp_1 = ceil_price(max(midt_2_proj_1 * (1. + self._edge_1 / 2.), ask0_1 - t1, bid0_1 + t1), t1)
    bidp_1 = floor_price(min(midt_2_proj_1 * (1. - self._edge_1 / 2.), ask0_1 - t1, bid0_1 + t1),
                         t1)

    midt_1_proj_2 = midt_1 - float(self._mid_basis_ma.average)
    askp_2 = ceil_price(max(midt_1_proj_2 * (1. + self._edge_2 / 2.), ask0_2 - t2, bid0_2 + t2), t2)
    bidp_2 = floor_price(min(midt_1_proj_2 * (1. - self._edge_2 / 2.), ask0_2 - t2, bid0_2 + t2),
                         t2)

    if not self.is_og_ready() or self._dry_run:
      return

    self.post_1(askp_1, bidp_1)
    self.post_2(askp_2, bidp_2)

    now_ts = get_timestamp()
    if now_ts - self._last_feature_dump_ts >= 2 * (10**9):
      self._last_feature_dump_ts = now_ts
      # "FEATURE_DUMP",time,time_hr,
      # ask0_1,bid0_1,askt_1,bidt_1,midt_1,
      # ask0_2,bid0_2,askt_2,bidt_2,midt_2,
      # mid_basis_ma,askp_1,bidp_1,askp_2,bidp_2
      self._logger.info(
          '\n"FEATURE_DUMP",%d,"%s",'
          '%f,%f,%f,%f,%f,'
          '%f,%f,%f,%f,%f,'
          '%f,%f,%f,%f,%f',
          book_ts,
          to_datetime(book_ts),
          ask0_1,
          bid0_1,
          askt_1,
          bidt_1,
          midt_1,
          ask0_2,
          bid0_2,
          askt_2,
          bidt_2,
          midt_2,
          float(self._mid_basis_ma.average),
          askp_1,
          bidp_1,
          askp_2,
          bidp_2)

  def post_1(self, askp, bidp):
    if self._sb_pos <= -self._imbalance_allowance:
      askp = None
    if self._bs_pos >= self._imbalance_allowance:
      bidp = None
    self._executor_1.manage_mm_orders(askp, bidp)

  def post_2(self, askp, bidp):
    if self._bs_pos <= -self._imbalance_allowance:
      askp = None
    if self._sb_pos >= self._imbalance_allowance:
      bidp = None
    self._executor_2.manage_mm_orders(askp, bidp)


def get_param_set(init_og_immediately=False):
  product_2 = BitmexFuturesProduct.FromStr('BTC-USD.PERPETUAL')
  product_1 = OkexFuturesProduct.FromStr('BTC-USD.THIS_WEEK')
  edge_bp_1 = 20.
  edge_bp_2 = 20.
  basis_ma_window_min = 10
  okex_og_config_name = 'okex_jshin2'
  bitmex_og_config_name = 'bitmex_xunke01'
  lot_size = 2000.

  strategy_params = {
      'product_1': product_1,
      'product_2': product_2,
      'basis_ma_window_min': basis_ma_window_min,
      'true_book_1': 10000.,
      'true_book_2': 10000.,
      'edge_bp_1': edge_bp_1,
      'edge_bp_2': edge_bp_2,
      'lot_size_1': lot_size,
      'lot_size_2': lot_size,
      'reserve': 20000.,
      'imbalance_allowance': lot_size * 2,
  }

  # Trading environment config
  if init_og_immediately:
    init_og_after_min = 0
  else:
    init_og_after_min = strategy_params['basis_ma_window_min']

  exit_after_min = init_og_after_min + 50
  env_config = {
      'first_currency': OkexCurrency.FromStrCurrency('BTC'),
      'first_product': product_1,
      'first_og_config_name': okex_og_config_name,
      'second_product': product_2,
      'second_og_config_name': bitmex_og_config_name,
      'init_order_gateway_after': datetime.timedelta(minutes=init_og_after_min),
      'exit_after': datetime.timedelta(minutes=exit_after_min)
  }

  return strategy_params, env_config


def main(argv):
  strategy_params, env_config = get_param_set(init_og_immediately=False)

  # Run strategy!
  strategy = PassivePassiveStrategy(strategy_params, dry_run=False)
  return StrategyBase.run_live(on_base_init_callback=strategy.on_base_init,
                               on_book_reset_callback=strategy.on_book_reset,
                               on_order_gateway_reset_callback=strategy.on_og_reset,
                               **env_config)


if __name__ == '__main__':
  logging.basicConfig(level='DEBUG', format='%(levelname)8s %(asctime)s %(name)s] %(message)s')
  sys.exit(main(sys.argv))
